from scipy.sparse import dia_matrix  # scipy.sparse：用于创建和操作稀疏矩阵。
from scipy.sparse.linalg import inv  # scipy.sparse.linalg：用于稀疏矩阵的线性代数操作，如求逆。
from numpy import pi  # numpy：用于数值计算。
import numpy as np
import matplotlib.pyplot as plt  # matplotlib.pyplot：用于绘图。

class FEM:
    def __init__(self, nodes, xmin=0, xmax=1):
        self.nodes = nodes   #节点数
        x = np.linspace(xmin, xmax, nodes)  # 定义了计算域的范围，生成等间距的节点位置
        self.x = x
        self.h = x[1] - x[0]  # 计算节点间距

    def Kmatrix(self):
        n = self.nodes #节点总数
        m = 1 / self.h * np.ones(n - 2)  # 排除边界节点，构造内部节点的矩阵，1 / self.h 是离散化后的系数。
        data = [m, -2 * m, m]  # 三对角矩阵的对角线元素，分别对应下对角线、主对角线和上对角线。
        offsets = [-1, 0, 1]  # 对角线的偏移量，-1 表示下对角线，0 表示主对角线，1 表示上对角线。
        K = dia_matrix((data, offsets), shape=(n - 2, n - 2)).tocsc()  # 创建稀疏矩阵并转换为压缩列存储格式CSC格式
        #shape = (n - 2, n - 2) 表示矩阵的大小，排除了边界节点。
        return K

    def bvec(self):
        return - np.sin(pi * self.x[1:-1]) * self.h  # 排除边界节点，构造内部节点的载荷向量

    # self.x[1:-1]：排除边界节点，取内部节点的位置。
    # np.sin(pi * self.x[1:-1])：计算载荷函数在内部节点的值。
    # *self.h：根据离散化公式，乘以节点间距。

    def solve(self):
        K = self.Kmatrix()  # 获取刚度矩阵
        b = self.bvec()  # 获取载荷向量
        phi = np.zeros(self.nodes)  # 初始化解向量，包含边界值
        phi[1:-1] = inv(K) * b  # 求解内部节点的值，使用稀疏矩阵的逆矩阵乘以载荷向量。
        return phi

    def compare(self):
        ground_truth = 1 / (pi ** 2) * np.sin(pi * self.x)  # 真实解
        fem_res = self.solve()  # 有限元解

        plt.plot(self.x, ground_truth, label="ground truth")  # 绘制真实解
        plt.plot(self.x, fem_res, label="finite element")  # 绘制有限元解
        plt.title("number of nodes = %s" % self.nodes)  # 设置标题
        plt.xlabel("x")  # 设置x轴标签
        plt.ylabel(r"$\phi(x)$")  # 设置y轴标签
        plt.legend(loc='best')  # 添加图例
        plt.show()  # 显示图像

# 创建FEM类的实例，节点数为101
fem = FEM(101)
# 调用compare方法，展示数值解与真实解的对比
fem.compare()